import os.path
import uuid
from contextvars import ContextVar
from time import sleep

import httpx
from celery import Celery
from fastapi import FastAPI, Depends, Request
from pydantic import BaseModel
from celery.result import AsyncResult
from .core.config import settings
from .core.exception_handler import register_exception

from .schemas.response import ResponseModel
from .tasks import video
from .tasks.celery_base import make_celery

from py_tools.logging import logger, setup_logging

app = FastAPI(debug=settings.DEBUG)
register_exception(app)

app_celery = make_celery()

setup_logging(log_dir="logs")

app = FastAPI()


def get_trace_id(request: Request):
    trace_id = request.headers.get("X-Trace-Id") or str(uuid.uuid4())
    request.state.trace_id = trace_id
    return trace_id


@app.middleware("http")
async def add_trace_id_header(request: Request, call_next):
    trace_id = get_trace_id(request)
    body = await request.body()  # 注意：body 读取后不能再次读取，后面处理要小心
    logger.info(f"Request: TraceID: {trace_id} | Body: {body.decode('utf-8')}")
    # 处理响应
    response = await call_next(request)
    # 记录响应日志（如果需要获取响应体，可以用自定义 Response 类包装 response）
    logger.info(f"Response: TraceID: {trace_id} | Status: {response.status_code}")
    response.headers["X-Trace-Id"] = trace_id
    return response


class Watermark(BaseModel):
    video_path: str


@app.post("/video/watermark")
async def video_watermark(watermark: Watermark, response_model=ResponseModel):
    result = video.watermark.apply_async(
        args=[watermark.video_path, os.path.join(settings.ASSETS_DIR, 'watermark.png'), (10, 10)]
    )
    return response_model(data=result.id)


@app.get("/task/{task_id}")
async def task_info(task_id: str, response_model=ResponseModel):
    # 获取任务结果对象
    result = AsyncResult(task_id, app=app_celery)

    if result.state == "STARTED":
        return response_model(data="任务排队中")
    if result.state == "PENDING":
        return response_model(data="任务未完成")
    if result.state == "FAILURE":
        return response_model(data="任务失败" + result.traceback)
    if result.state == "SUCCESS":
        return response_model(data=result.info)


@app.get("/callback")
async def callback(task_id: str):
    async with httpx.AsyncClient() as client:
        # 例如回调第三方接口
        await client.post("http://localhost:8001/test", timeout=5)


@app.get("/test")
async def test(response_model=ResponseModel):
    logger.info("访问根路径")
    return response_model(data="success")
